# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: leon

import os
import datetime
import ujson as json
import json as base_json
import functools

from coin.base.book.types import TradeSide, create_market_trade_event
from coin.exchange.base.ws.handler import WebsocketHandler
from coin.exchange.huobi.ws.subscriber import parse_channel
from coin.exchange.huobi_futures.book_builder import BookBuilder
from coin.exchange.huobi_futures.kr_rest.constants import native_subscription_symbol_list
from coin.exchange.huobi_futures.kr_rest.futures_product import (
    HuobiFuturesProduct,)
import coin.exchange.huobi_futures.kr_rest.contract_util as contract_util
from coin.exchange.base.ws.trade_aggregator import TradeAggregator


@functools.lru_cache(maxsize=128)
def _get_product_from_channel_name(channel, feed_type, as_of_timestamp):
  native_subscription_symbol, _feed_type = parse_channel(channel)
  assert native_subscription_symbol in native_subscription_symbol_list
  assert _feed_type == feed_type
  base_str, native_subscription_contract = native_subscription_symbol.split('_')
  currency_pair = "%s_%s" % (base_str, 'USD')
  contract_type = contract_util.to_contract_type(native_subscription_contract)
  current_time = datetime.datetime.fromtimestamp(as_of_timestamp / 1.e9)
  return HuobiFuturesProduct.FromCurrencyPairAndContractType(currency_pair,
                                                             contract_type,
                                                             current_time)


def get_product_from_channel_name(channel, feed_type, as_of_timestamp):
  as_of_timestamp = (as_of_timestamp // (1800 * 10**9)) * (1800 * 10**9)
  return _get_product_from_channel_name(channel, feed_type, as_of_timestamp)


class HuobiFuturesHandler(WebsocketHandler):
  def __init__(self,
               *,
               event_dist=None,
               feed_checker_config=None,
               agg_view=True,
               book_builder_opts=None):
    super().__init__(event_dist=event_dist, feed_checker_config=feed_checker_config)
    self._trade_processor = TradeAggregator(handler=self, agg_view=agg_view)
    self._book_builder_opts = book_builder_opts or {}

  def on_depth(self, record, queue_data, topic_data):
    data = record.data
    if not isinstance(data, dict) and not isinstance(data, list):
      data = json.loads(data.tobytes())

    product = get_product_from_channel_name(data['ch'], 'depth', record.timestamp)
    native_symbol = product.native_symbol
    if not self.is_subscribed('book', native_symbol):
      return

    if native_symbol in self._books:
      book = self._books[native_symbol]
      book.clear_book()
    else:
      current_datetime = datetime.datetime.fromtimestamp(record.timestamp / 1.e9)
      book = BookBuilder(native_symbol,
                         current_datetime,
                         book_builder_observer=self._book_builder_observer,
                         **self._book_builder_opts)
    assert 'tick' in data
    try:
      bids = data['tick']['bids']
      asks = data['tick']['asks']
    except KeyError:
      return
    book.snapshot(bids, asks)
    book.update_timestamp(record.timestamp)
    book.update_topic_id(record.topic_id)
    self._books[native_symbol] = book
    self.ready = True
    if self._feed_checker:
      self._feed_checker.on_book(product=product, book=book, timestamp=record.timestamp)
    self.publish_event('book', book.symbol, book)

  def on_trade(self, record, queue_data, topic_data):
    data = record.data
    if not isinstance(data, dict) and not isinstance(data, list):
      try:
        data = json.loads(data.tobytes())
      except ValueError:
        data = base_json.loads(data.tobytes())

    product = get_product_from_channel_name(data['ch'], 'trade', record.timestamp)
    native_symbol = product.native_symbol
    if not self.is_subscribed('trade', native_symbol):
      return

    self._trade_processor.on_begin_trade()

    data = data['tick']
    for data in data['data']:
      price = float(data['price'])
      # Huobi Futures sends double-counted qty
      qty = float(data['amount']) / 2
      if data['direction'] == 'buy':
        side = TradeSide.TRADE_BUY_SIDE
      elif data['direction'] == 'sell':
        side = TradeSide.TRADE_SELL_SIDE
      else:
        raise ValueError('Unknown side: %s', data)
      trade_event = create_market_trade_event(record.timestamp,
                                              native_symbol,
                                              price,
                                              qty,
                                              side,
                                              topic_id=record.topic_id)
      self._trade_processor.on_mid_trade(trade_event)
      if self._feed_checker:
        self._feed_checker.on_trade(product=product,
                                    trade_event=trade_event,
                                    timestamp=record.timestamp)

    self._trade_processor.on_end_trade()


if ('RUST_FEED_HANDLER' in os.environ) and os.environ['RUST_FEED_HANDLER'] == '1':
  try:
    from native_rust import feed_handler
    print("huobi futures: using rust feed handler")

    class HuobiFuturesRustHandler(HuobiFuturesHandler):
      def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._rust_handler = feed_handler.HuobiFuturesHandler()
        self._rust_handler.register_book_callback(functools.partial(self.publish_event, 'book'))

      def post_subscribe(self, event_type, symbol, callback):
        assert type(symbol) == str
        if event_type == 'book':
          self._rust_handler.register_book_symbol(symbol)

      def on_depth(self, record, queue_data, topic_data):
        if isinstance(record.data, (dict, list)):
          record_bytes = json.dumps(record.data).encode()
        else:
          record_bytes = record.data.tobytes()

        self._rust_handler.on_snapshot(record.timestamp, record_bytes, {}, {})

    HuobiFuturesHandler = HuobiFuturesRustHandler

  except ImportError:
    import logging
    logging.warning('Unable to import rust feed_handler.')
